package com.ustc.hewei.detectbeats.detect;

import com.alibaba.fastjson2.JSON;
import com.ustc.hewei.rpc.client.net.SendHandler;
import com.ustc.hewei.rpc.common.Utils.MessageSerialize;
import com.ustc.hewei.rpc.common.protocol.*;
import io.netty.bootstrap.Bootstrap;
import io.netty.channel.ChannelInitializer;
import io.netty.channel.ChannelOption;
import io.netty.channel.FixedRecvByteBufAllocator;
import io.netty.channel.nio.NioEventLoopGroup;
import io.netty.channel.socket.SocketChannel;
import io.netty.channel.socket.nio.NioSocketChannel;
import lombok.extern.slf4j.Slf4j;
import org.springframework.data.redis.core.Cursor;
import org.springframework.data.redis.core.RedisCallback;
import org.springframework.data.redis.core.ScanOptions;
import org.springframework.data.redis.core.StringRedisTemplate;

import javax.annotation.Resource;
import java.io.IOException;
import java.util.*;
import java.util.stream.Collectors;

/**
* @description: 发送心跳检测信号
* @author hewei
* @date 2022/8/12 11:00
* @version 1.0
*/

@Slf4j
public class SendBeats {
    @Resource
    private StringRedisTemplate stringRedisTemplate;

    // 从bean工厂中引入
    private String protocol;

    public SendBeats(String protocol) {
        this.protocol = protocol;
    }

    /*
     * @description: 从redis注册中心中拉取所有的服务，采用模糊查询的方式
     * @return java.util.HashMap<java.lang.String,java.util.List<java.lang.String>>
                键为服务的完整的ip地址，值为这个地址下所有的服务名称所对应的redis下的键
     * @author: hewei
     * @date: 2022/8/13 10:24
     **/
    public HashMap<String, List<String>> findService() {
        // 模糊查询的条件
        String pattern = "rpc/*";
        HashMap<String, List<String>> allServices = new HashMap<>();
        // 执行模糊查询，采用scan的方式（不阻塞），使用keys的方式会阻塞，因为redis是单线程
        Set<String> keySet = stringRedisTemplate.execute((RedisCallback<Set<String>>) connection -> {
            Set<String> keysTemp = new HashSet<>();
            Cursor<byte[]> cursor = connection.scan(ScanOptions.scanOptions().match(pattern).count(10000).build());
            while (cursor.hasNext()) {
                keysTemp.add(new String(cursor.next()));
            }
            return keysTemp;
        });
        // 没找到rpc/开头的键值，代表没有服务注册
        if (keySet == null || keySet.size() == 0) {
            log.info("注册中心暂无服务！");
            return allServices;
        }
        // 查询每个键对应的set集合，然后按照ip进行归类
        for (String key : keySet) {
            List<Service> serviceList = stringRedisTemplate.opsForSet().members(key).stream()
                    .map(s -> JSON.parseObject(s, Service.class))
                    .collect(Collectors.toList());
            for (Service service : serviceList) {
                if (!allServices.containsKey(service.getAddress())) allServices.put(service.getAddress(), new ArrayList<>());
                allServices.get(service.getAddress()).add(key);
            }
        }
        return allServices;
    }

    /*
     * @description: 发送心跳的主方法
     * @author: hewei
     * @date: 2022/8/13 10:26
     **/
    public void send() throws IOException {
        HashMap<String, List<String>> allServiceList = findService();
        for (Map.Entry<String, List<String>> entry : allServiceList.entrySet()) {
            String ip = entry.getKey();
            List<String> keys = entry.getValue();
            RequestWarp requestWarp = new RequestWarp();
            requestWarp.setRequestType(RequestType.KEEP_ALIVE);
            byte[] data = MessageSerialize.objectToByteArray(requestWarp);
            HashMap<String, String> removedService = getRemovedService(keys, ip);
            try {
                byte[] response = sendRequest(data, ip);
                ResponseWarp responseWarp = MessageSerialize.byteArrayToResp(response);
                if (responseWarp.getStatus() != Status.KEEP_ALIVE) {
                    for (Map.Entry<String, String> removeEntry : removedService.entrySet()) {
                        stringRedisTemplate.opsForSet().remove(removeEntry.getKey(), removeEntry.getValue());
                    }
                    log.info("已删除服务 {} 下所有服务!", ip);
                }
            } catch (Exception e) {
                log.error("心跳检测异常！");
                for (Map.Entry<String, String> removeEntry : removedService.entrySet()) {
                    stringRedisTemplate.opsForSet().remove(removeEntry.getKey(), removeEntry.getValue());
                }
                log.info("已删除服务 {} 下所有服务!", ip);
            }
        }
    }

    /*
     * @description: 获取到要删除的ip对应的服务列表
     * @param keys 要删除的ip地址下所有服务在redis中的键
     * @param ip 要删除的ip地址
     * @return java.util.HashMap<java.lang.String,java.lang.String>
                键为所要在redis中删除的键，值为封装的Service对象序列化之后的字符串
     * @author: hewei
     * @date: 2022/8/13 10:28
     **/
    public HashMap<String, String> getRemovedService(List<String> keys, String ip) {
        HashMap<String, String> map = new HashMap<>();
        for (String key : keys) {
            Service service = new Service();
            // 将redis中的键进行分割，得到服务的名称
            String serviceName = key.substring(key.indexOf("/") + 1, key.lastIndexOf("/"));
            service.setAddress(ip).setName(serviceName).setProtocol(protocol);
            String jsonString = JSON.toJSONString(service);
            map.put(key, jsonString);
        }
        return map;
    }

    /*
     * @description: 向服务端发送心跳信号
     * @param data 发送的RequestWarp对象序列化之后的字节数组
     * @param fullIp 服务端的全地址
     * @return byte[] 返回的ResponseWarp
     * @author: hewei
     * @date: 2022/8/13 10:32
     **/
    public byte[] sendRequest(byte[] data, String fullIp) {
        SendHandler sendHandler = new SendHandler(data);
        byte[] response;

        //将地址进行分割，因为形式是"ip：port"
        String[] split = fullIp.split(":");
        int port = Integer.parseInt(split[1]);
        String ip = split[0];

        NioEventLoopGroup group = new NioEventLoopGroup();
        try {
            new Bootstrap().group(group).channel(NioSocketChannel.class)
                    .option(ChannelOption.TCP_NODELAY, true)
                    //设置可以接收的大小，默认为2048个字节，不够用
                    .option(ChannelOption.RCVBUF_ALLOCATOR, new FixedRecvByteBufAllocator(65535))
                    .handler(new ChannelInitializer<SocketChannel>() {
                        @Override
                        protected void initChannel(SocketChannel socketChannel) throws Exception {
                            //调用自定义的消息处理器
                            socketChannel.pipeline().addLast(sendHandler);
                        }
                    }).connect(ip, port).sync();
            response = (byte[]) sendHandler.getResponse();
            log.info("从 {} 获取心跳成功！", fullIp);
        } catch (InterruptedException e) {
            throw new RuntimeException("心跳检测网络连接异常！", e);
        } finally {
            group.shutdownGracefully();
        }
        return response;
    }
}
